from jqdata import *

"""
|以下代码是jq回测使用
# 部分方法来自：https://www.joinquant.com/post/1401
# 部分来自聚宽文章：https://www.joinquant.com/post/1957
# peg选股+海龟择时策略
# 2012-01-01 到 2016-03-10, ￥1000000, 分钟
作者：Huchian
================================================================================
"""
import datetime
import pandas as pd
# import sys
# enable_profile()


'''
================================================================================
总体回测前
================================================================================
'''
#总体回测前要做的事情
def initialize(context):
    set_params()                             # 设置策略常量
    set_variables()                          # 设置中间变量
    set_backtest()                           # 设置回测条件

    g.ts = TurtleSystem()
    """ 设置两个子账户 """
    sys1_init_cash = context.portfolio.starting_cash * g.ts.sys1_ratio
    sys2_init_cash = context.portfolio.starting_cash - sys1_init_cash
    set_subportfolios([SubPortfolioConfig(cash=sys1_init_cash, type='stock'),\
                    SubPortfolioConfig(cash=sys2_init_cash, type='stock')])
    

#设置策略参数
def set_params():
    g.tc = 15                                # 调仓天数
    g.num_stocks = 60                        # 每次调仓选取的最大股票数量

#设置中间变量
def set_variables():
    g.t = 0                                  # 记录回测运行的天数
    g.if_trade = False                       # 当天是否交易


#设置回测条件
def set_backtest():
    # 作为判断策略好坏和一系列风险值计算的基准
    # 设定沪深300作为基准
    set_benchmark('000300.XSHG')

    set_option('use_real_price',True)        # 用真实价格交易
    log.set_level('order','warning')           # 设置报错等级

'''
================================================================================
每天开盘前
================================================================================
'''
def before_trading_start(context):
    set_slip_fee(context) 
    g.orders_unfinished = {}

    if g.t%g.tc==0:
        g.if_trade=True                          # 每g.tc天，调仓一次
        set_slip_fee(context)                    # 设置手续费与手续费
        g.stocks=get_index_stocks('000300.XSHG') # 设置沪深300为初始股票池
        
    g.t+=1
    # 待买入的g.num_stocks支股票，list类型
    list_to_buy = stocks_to_buy(g.stocks, g.num_stocks)
    # list_to_buy = stocks_to_buy4(context, g.num_stocks)

    # 设置可行股票池
    list_to_buy = set_feasible_stocks(list_to_buy)
    # 读取行情
    days = g.ts.get_max_prices_days()
    securities_data_history = {}
    for security in list_to_buy:
        prices = attribute_history(security, days, '1d',('high','low','close','pre_close'))
        #过滤价格太高的股票, 1股超过250 4手就超过10万 ,预留一个涨停板10% TODO 可配置
        if prices['close'][-1] > 227: 
            continue
        
        #涨停也无法最小突破的 TODO
        if prices['close'][-1] * 1.1 < max(prices['high'][-20:]):
            continue

        sec_info = get_security_info(security)

        securities_data_history[security] = {}
        securities_data_history[security]['prices'] = prices
        securities_data_history[security]['end_date'] = sec_info.end_date
        securities_data_history[security]['type'] = sec_info.type
        securities_data_history[security]['in_buy_list'] = True
    
    # 需要更新行情的
    securities_info_to_refresh = g.ts.securities_info_to_refresh()
    for security in securities_info_to_refresh:
        securities_data_history[security] = {}
        securities_data_history[security]['prices'] = attribute_history(security, days, '1d',('high','low','close','pre_close'))
    
    g.ts.refresh_basic_data(context.portfolio, securities_data_history)
    
    
# 设置可行股票池：过滤掉当日停牌的股票
# 输入：initial_stocks为list类型,表示初始股票池； context（见API）
# 输出：unsuspened_stocks为list类型，表示当日未停牌的股票池，即：可行股票池
def set_feasible_stocks(initial_stocks):
    # 判断初始股票池的股票是否停牌，返回list
    paused_info = []
    current_data = get_current_data()
    for i in initial_stocks:
        paused_info.append(current_data[i].paused)
    df_paused_info = pd.DataFrame({'paused_info':paused_info},index = initial_stocks)
    unsuspened_stocks =list(df_paused_info.index[df_paused_info.paused_info == False])
    return unsuspened_stocks


# 根据不同的时间段设置滑点与手续费
def set_slip_fee(context):
    # 将滑点设置为0
    set_slippage(FixedSlippage(0)) 
    # 根据不同的时间段设置手续费
    dt=context.current_dt
    
    if dt>datetime.datetime(2013,1, 1):
        set_commission(PerTrade(buy_cost=0.0003, sell_cost=0.0013, min_cost=5)) 
        
    elif dt>datetime.datetime(2011,1, 1):
        set_commission(PerTrade(buy_cost=0.001, sell_cost=0.002, min_cost=5))
            
    elif dt>datetime.datetime(2009,1, 1):
        set_commission(PerTrade(buy_cost=0.002, sell_cost=0.003, min_cost=5))
                
    else:
        set_commission(PerTrade(buy_cost=0.003, sell_cost=0.004, min_cost=5))


'''
================================================================================
每天交易时
================================================================================
'''
# 按分钟回测
def handle_data(context, data):
    """ 处理9:30的feature：data为前一交易日 """
    dt = context.current_dt
    if dt.hour == 9 and dt.minute == 30:
        return None

    # 执行委托
    """
    [sec, -position, 'sell_for_cash', sec_info_data[sec]['current_price'], 'long', self.pindex]
    """
    orders = g.ts.trade_strategy(context.portfolio, data, 'price')

    place_orders(orders)
    check_unfinished_orders()

    """ 最后1分钟，处理现金 """
    if dt.strftime('%Y-%m-%d') >= '2013-04-18' and dt.hour == 14 and dt.minute == 59:
        orders = g.ts.manage_cash(context.portfolio, data, 'price')
        place_orders(orders)
    return None


def place_orders(orders):

    """ 订单委托 """
    for o in orders:
        sec = o[0]
        position = o[1]
        order_type = o[2]
        side = o[4]
        pindex = o[5]
        """
        UserOrder({'order_id': 1648647936, 'security': '601898.XSHG', 'amount': 3200, 'filled': 3200, 'price': 8.03, 'status': held, 'add_time': datetime.datetime(2022, 3, 7, 9, 31), 'side': 'long', 'action': 'open', 'pindex': 1, 'style': MarketOrderStyle: _limit_price=0.0})
        """
        r = order(sec, position, pindex=pindex, side=side)
        if r is not None:
            # 订单执行成功
            if r.status == OrderStatus.held:
                g.ts.after_order_success(r, order_type)
            else:
                # 简单处理 TODO 处理所有状态
                g.orders_unfinished[r.order_id] = o
        else:
            """ 如需要详细原因调整 log.order 级别到warning """
            log.warning("委托不成功")


def check_unfinished_orders():
    """ 检查之前未完成的是否完成了 """

    orders_unfinished = {}
    for order_id in g.orders_unfinished:
        orders =  get_orders(order_id=order_id)
        order_type = g.orders_unfinished[order_id][2]

        if orders[order_id].status == OrderStatus.held: # 订单执行成功
            g.ts.after_order_success(orders[order_id], order_type)
        else:                                           # 订单执行不成功
            orders_unfinished[order_id] = g.orders_unfinished[order_id]
    g.orders_unfinished = orders_unfinished
    pass

# 计算股票的PEG值
# 输入：context(见API)；stock_list为list类型，表示股票池
# 输出：df_PEG为dataframe: index为股票代码，data为相应的PEG值
def get_PEG(stock_list): 
    # 查询股票池里股票的市盈率，收益增长率
    q_PE_G = query(valuation.code, valuation.pe_ratio, indicator.inc_net_profit_year_on_year
                 ).filter(valuation.code.in_(stock_list)) 
    # 得到一个dataframe：包含股票代码、市盈率PE、收益增长率G
    # 默认date = context.current_dt的前一天,使用默认值，避免未来函数，不建议修改
    df_PE_G = get_fundamentals(q_PE_G)
    # 筛选出成长股：删除市盈率或收益增长率为负值的股票
    df_Growth_PE_G = df_PE_G[(df_PE_G.pe_ratio >0)&(df_PE_G.inc_net_profit_year_on_year >0)]
    # 去除PE或G值为非数字的股票所在行
    df_Growth_PE_G.dropna()
    # 得到一个Series：存放股票的市盈率TTM，即PE值
    Series_PE = df_Growth_PE_G.ix[:,'pe_ratio']
    # 得到一个Series：存放股票的收益增长率，即G值
    Series_G = df_Growth_PE_G.ix[:,'inc_net_profit_year_on_year']
    # 得到一个Series：存放股票的PEG值
    Series_PEG = Series_PE/Series_G
    # 将股票与其PEG值对应
    Series_PEG.index = df_Growth_PE_G.ix[:,0]
    # 将Series类型转换成dataframe类型
    df_PEG = pd.DataFrame(Series_PEG)
    return df_PEG
    
# 获得买入信号
# 输入：context(见API)
# 输出：list_to_buy为list类型,表示待买入的g.num_stocks支股票
def stocks_to_buy(stocks = [], num = 10):
    list_to_buy = []
    # 得到一个dataframe：index为股票代码，data为相应的PEG值
    df_PEG = get_PEG(stocks)
    # 将股票按PEG升序排列，返回daraframe类型
    try:
        df_sort_PEG = df_PEG.sort(columns=[0], ascending=[1])
    except AttributeError:
        df_sort_PEG = df_PEG.sort_values(by=[0], ascending=[1])
    # 将存储有序股票代码index转换成list并取前g.num_stocks个为待买入的股票，返回list
    for i in range(num):
        if df_sort_PEG.ix[i,0] < 0.5:
            list_to_buy.append(df_sort_PEG.index[i])
    return list_to_buy


'''
================================================================================
每天收盘后
================================================================================
'''
def after_trading_end(context):
    orders = get_open_orders()
    for _order in orders.values():
        print(_order)
    
    g.ts.position_check(context.portfolio)


def on_event(context, event):
    """ 
    建议用户使用isinstance对事件类型进行判断。 目前已支持的事件有：
    DividendsEvent：分红送股事件
    ForcedLiquidationEvent：强行平仓事件
    """
    pass

def on_strategy_end(context):
    """
    策略运行结束后，输出一些内部统计结果 TODO
    """
    for subportfolio in context.portfolio.subportfolios:
        print(subportfolio)

################################################################################

# 
from securities_info import *
# 此行需要替换为文件 turtle_system.py 中的内容
from turtle_system import *

